# 导入必要的库
import json
import tempfile

import numpy as np
import copy
import time
import torch
import torch._six

from pycocotools.cocoeval import COCOeval
from pycocotools.coco import COCO
import pycocotools.mask as mask_util

from collections import defaultdict

import utility.utils as utils


# COCO评估器类，用于评估目标检测和分割模型的性能
class CocoEvaluator(object):
    """
    COCO评估器类，用于评估目标检测和分割模型的性能。
    """
    def __init__(self, coco_gt, iou_types):
        """
        初始化CocoEvaluator。

        :param coco_gt: COCO格式的真实标注数据
        :param iou_types: 评估的IoU类型，如 'bbox' 或 'segm'
        """
        # 确保iou_types是列表或元组
        assert isinstance(iou_types, (list, tuple))
        # 深拷贝真实标注数据，避免修改原始数据
        coco_gt = copy.deepcopy(coco_gt)
        self.coco_gt = coco_gt

        # 存储iou类型
        self.iou_types = iou_types
        # 初始化COCO评估对象字典
        self.coco_eval = {}
        for iou_type in iou_types:
            # 为每种iou类型创建一个COCOeval对象
            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)

        # 初始化图像ID列表
        self.img_ids = []
        # 初始化评估图像结果字典
        self.eval_imgs = {k: [] for k in iou_types}

    def update(self, predictions):
        """
        更新评估器的预测结果。

        :param predictions: 模型的预测结果
        """
        # 获取预测结果中的图像ID
        img_ids = list(np.unique(list(predictions.keys())))
        self.img_ids.extend(img_ids)

        for iou_type in self.iou_types:
            # 根据iou类型准备预测结果
            results = self.prepare(predictions, iou_type)
            # 如果有预测结果，加载结果；否则创建一个空的COCO对象
            coco_dt = loadRes(self.coco_gt, results) if results else COCO()
            coco_eval = self.coco_eval[iou_type]

            coco_eval.cocoDt = coco_dt
            coco_eval.params.imgIds = list(img_ids)
            # 对每张图像进行评估
            img_ids, eval_imgs = evaluate(coco_eval)

            self.eval_imgs[iou_type].append(eval_imgs)

    def synchronize_between_processes(self):
        """
        在多进程之间同步评估结果。
        """
        for iou_type in self.iou_types:
            # 合并多进程的评估结果
            self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
            # 创建通用的COCO评估对象
            create_common_coco_eval(self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])

    def accumulate(self):
        """
        累积评估结果。
        """
        for coco_eval in self.coco_eval.values():
            coco_eval.accumulate()

    def summarize(self):
        """
        总结评估结果并打印。
        """
        for iou_type, coco_eval in self.coco_eval.items():
            print("IoU metric: {}".format(iou_type))
            coco_eval.summarize()

    def prepare(self, predictions, iou_type):
        """
        根据iou类型准备预测结果。

        :param predictions: 模型的预测结果
        :param iou_type: IoU类型
        :return: 准备好的预测结果
        """
        if iou_type == "bbox":
            return self.prepare_for_coco_detection(predictions)
        elif iou_type == "segm":
            return self.prepare_for_coco_segmentation(predictions)
        elif iou_type == "keypoints":
            return self.prepare_for_coco_keypoint(predictions)
        else:
            raise ValueError("Unknown iou type {}".format(iou_type))

    def prepare_for_coco_detection(self, predictions):
        """
        准备目标检测的预测结果。

        :param predictions: 模型的预测结果
        :return: COCO格式的预测结果
        """
        coco_results = []
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            boxes = prediction["boxes"]
            boxes = convert_to_xywh(boxes).tolist()
            scores = prediction["scores"].tolist()
            labels = prediction["labels"].tolist()

            coco_results.extend(
                [
                    {
                        "image_id": original_id,
                        "category_id": labels[k],
                        "bbox": box,
                        "score": scores[k],
                    }
                    for k, box in enumerate(boxes)
                ]
            )
        return coco_results

    def prepare_for_coco_segmentation(self, predictions):
        """
        准备分割的预测结果。

        :param predictions: 模型的预测结果
        :return: COCO格式的预测结果
        """
        coco_results = []
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            scores = prediction["scores"]
            labels = prediction["labels"]
            masks = prediction["masks"]

            masks = masks > 0.5

            scores = prediction["scores"].tolist()
            labels = prediction["labels"].tolist()

            rles = [
                mask_util.encode(np.array(mask[0, :, :, np.newaxis], order="F"))[0]
                for mask in masks
            ]
            for rle in rles:
                rle["counts"] = rle["counts"].decode("utf-8")

            coco_results.extend(
                [
                    {
                        "image_id": original_id,
                        "category_id": labels[k],
                        "segmentation": rle,
                        "score": scores[k],
                    }
                    for k, rle in enumerate(rles)
                ]
            )
        return coco_results

    def prepare_for_coco_keypoint(self, predictions):
        """
        准备关键点检测的预测结果。

        :param predictions: 模型的预测结果
        :return: COCO格式的预测结果
        """
        coco_results = []
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            boxes = prediction["boxes"]
            boxes = convert_to_xywh(boxes).tolist()
            scores = prediction["scores"].tolist()
            labels = prediction["labels"].tolist()
            keypoints = prediction["keypoints"]
            keypoints = keypoints.flatten(start_dim=1).tolist()

            coco_results.extend(
                [
                    {
                        "image_id": original_id,
                        "category_id": labels[k],
                        'keypoints': keypoint,
                        "score": scores[k],
                    }
                    for k, keypoint in enumerate(keypoints)
                ]
            )
        return coco_results


# 将边界框从 [xmin, ymin, xmax, ymax] 转换为 [xmin, ymin, width, height]
def convert_to_xywh(boxes):
    xmin, ymin, xmax, ymax = boxes.unbind(1)
    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)


# 合并多进程的评估结果
def merge(img_ids, eval_imgs):
    """
    合并多进程的评估结果。

    :param img_ids: 图像ID列表
    :param eval_imgs: 评估图像结果
    :return: 合并后的图像ID和评估结果
    """
    all_img_ids = utils.all_gather(img_ids)
    all_eval_imgs = utils.all_gather(eval_imgs)

    merged_img_ids = []
    for p in all_img_ids:
        merged_img_ids.extend(p)

    merged_eval_imgs = []
    for p in all_eval_imgs:
        merged_eval_imgs.append(p)

    merged_img_ids = np.array(merged_img_ids)
    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)

    # 保留唯一的图像ID
    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
    merged_eval_imgs = merged_eval_imgs[..., idx]

    return merged_img_ids, merged_eval_imgs


# 创建一个通用的COCO评估对象
def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
    """
    创建一个通用的COCO评估对象。

    :param coco_eval: COCO评估对象
    :param img_ids: 图像ID
    :param eval_imgs: 评估图像结果
    """
    img_ids, eval_imgs = merge(img_ids, eval_imgs)
    img_ids = list(img_ids)
    eval_imgs = list(eval_imgs.flatten())

    coco_eval.evalImgs = eval_imgs
    coco_eval.params.imgIds = img_ids
    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)


# 从pycocotools中复制的代码，移除了打印语句并修复了Python3中关于unicode未定义的错误
def createIndex(self):
    """
    创建索引，用于快速查找图像、标注和类别信息。
    """
    # 创建索引
    anns, cats, imgs = {}, {}, {}
    imgToAnns, catToImgs = defaultdict(list), defaultdict(list)
    if 'annotations' in self.dataset:
        for ann in self.dataset['annotations']:
            imgToAnns[ann['image_id']].append(ann)
            anns[ann['id']] = ann

    if 'images' in self.dataset:
        for img in self.dataset['images']:
            imgs[img['id']] = img

    if 'categories' in self.dataset:
        for cat in self.dataset['categories']:
            cats[cat['id']] = cat

    if 'annotations' in self.dataset and 'categories' in self.dataset:
        for ann in self.dataset['annotations']:
            catToImgs[ann['category_id']].append(ann['image_id'])

    # 创建类成员
    self.anns = anns
    self.imgToAnns = imgToAnns
    self.catToImgs = catToImgs
    self.imgs = imgs
    self.cats = cats


maskUtils = mask_util


# 加载结果文件并返回一个结果API对象
def loadRes(self, resFile):
    """
    加载结果文件并返回一个结果API对象。

    :param resFile: 结果文件名
    :return: 结果API对象
    """
    res = COCO()
    res.dataset['images'] = [img for img in self.dataset['images']]

    if isinstance(resFile, torch._six.string_classes):
        anns = json.load(open(resFile))
    elif type(resFile) == np.ndarray:
        anns = self.loadNumpyAnnotations(resFile)
    else:
        anns = resFile
    assert type(anns) == list, 'results in not an array of objects'
    annsImgIds = [ann['image_id'] for ann in anns]
    assert set(annsImgIds) == (set(annsImgIds) & set(self.getImgIds())), \
        'Results do not correspond to current coco set'
    if 'caption' in anns[0]:
        imgIds = set([img['id'] for img in res.dataset['images']]) & set([ann['image_id'] for ann in anns])
        res.dataset['images'] = [img for img in res.dataset['images'] if img['id'] in imgIds]
        for id, ann in enumerate(anns):
            ann['id'] = id + 1
    elif 'bbox' in anns[0] and not anns[0]['bbox'] == []:
        res.dataset['categories'] = copy.deepcopy(self.dataset['categories'])
        for id, ann in enumerate(anns):
            bb = ann['bbox']
            x1, x2, y1, y2 = [bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]]
            if 'segmentation' not in ann:
                ann['segmentation'] = [[x1, y1, x1, y2, x2, y2, x2, y1]]
            ann['area'] = bb[2] * bb[3]
            ann['id'] = id + 1
            ann['iscrowd'] = 0
    elif 'segmentation' in anns[0]:
        res.dataset['categories'] = copy.deepcopy(self.dataset['categories'])
        for id, ann in enumerate(anns):
            # 现在只支持压缩RLE格式作为分割结果
            ann['area'] = maskUtils.area(ann['segmentation'])
            if 'bbox' not in ann:
                ann['bbox'] = maskUtils.toBbox(ann['segmentation'])
            ann['id'] = id + 1
            ann['iscrowd'] = 0
    elif 'keypoints' in anns[0]:
        res.dataset['categories'] = copy.deepcopy(self.dataset['categories'])
        for id, ann in enumerate(anns):
            s = ann['keypoints']
            x = s[0::3]
            y = s[1::3]
            x0, x1, y0, y1 = np.min(x), np.max(x), np.min(y), np.max(y)
            ann['area'] = (x1 - x0) * (y1 - y0)
            ann['id'] = id + 1
            ann['bbox'] = [x0, y0, x1 - x0, y1 - y0]
    res.dataset['annotations'] = anns
    createIndex(res)
    return res


# 在给定的图像上运行每图像评估，并将结果存储在self.evalImgs中
def evaluate(self):
    """
    在给定的图像上运行每图像评估，并将结果存储在self.evalImgs中。

    :return: None
    """
    p = self.params
    # 如果useSegm在params中指定，则添加向后兼容性
    if p.useSegm is not None:
        p.iouType = 'segm' if p.useSegm == 1 else 'bbox'
        print('useSegm (已弃用) 不为None。运行 {} 评估'.format(p.iouType))
    # print('评估注释类型 *{}*'.format(p.iouType))
    p.imgIds = list(np.unique(p.imgIds))
    if p.useCats:
        p.catIds = list(np.unique(p.catIds))
    p.maxDets = sorted(p.maxDets)
    self.params = p

    self._prepare()
    # 遍历所有图像、类别和面积范围，计算评估结果
    catIds = p.catIds if p.useCats else [-1]

    if p.iouType == 'segm' or p.iouType == 'bbox':
        computeIoU = self.computeIoU
    elif p.iouType == 'keypoints':
        computeIoU = self.computeOks
    self.ious = {
        (imgId, catId): computeIoU(imgId, catId)
        for imgId in p.imgIds
        for catId in catIds}

    evaluateImg = self.evaluateImg
    maxDet = p.maxDets[-1]
    evalImgs = [
        evaluateImg(imgId, catId, areaRng, maxDet)
        for catId in catIds
        for areaRng in p.areaRng
        for imgId in p.imgIds
    ]
    # 合并多进程的评估结果
    evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds))
    self._paramsEval = copy.deepcopy(self.params)
    # toc = time.time()
    # print('DONE (t={:0.2f}s).'.format(toc-tic))
    return p.imgIds, evalImgs

# 直接从pycocotools中复制的代码结束，移除了打印语句